package swrr

import (
	"context"
	"fmt"
	"gitee.com/xiao_hange/go-admin-pkg/pkg/logger"
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/base"
	"strconv"
	"strings"
	"sync"
)

// 平滑加权负载均衡算法 (Smooth Weighted Round Robin) swrr
// 在本swrr上添加了节点过滤功能不同用户可以选择不同身份的节点
const name = "go_admin_swrr"
const defaultLabelType = "ordinary"

func init() {
	balancer.Register(base.NewBalancerBuilder(
		name,
		&PickerBuilder{
			l: logger.NewGrpcLogger(),
		},
		base.Config{HealthCheck: true},
	))
}

// PickerBuilder 实现 PickerBuilder  google.golang.org/grpc/balancer/base 接口
type PickerBuilder struct {
	SomeInfo string
	l        logger.GrpcLogger
}

// Build
// ReadySCs map[balancer.SubConn]SubConnInfo
// sc => SubConn
// sc => SubConnInfo

func (p *PickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
	//conns := make([]*conn, 0, len(info.ReadySCs))
	conns := make(map[string][]*conn)
	for sc, sci := range info.ReadySCs {
		cc := &conn{
			cc:   sc,
			addr: sci.Address.Addr,
		}
		md, ok := sci.Address.Metadata.(map[string]string)

		if ok {
			weightVal := md["weight"]
			cc.weight, _ = strconv.Atoi(weightVal)
			if lb, ok := md["label"]; ok {
				cc.lb = strings.ToLower(lb)
			}
		}
		if cc.weight == 0 {
			cc.weight = 10
		}
		if cc.lb == "" {
			cc.lb = defaultLabelType
		}
		cc.currentWeight = cc.weight
		//conns = append(conns, cc)
		conns[cc.lb] = append(conns[cc.lb], cc)
	}
	return &Picker{
		conns: conns,
		l:     p.l,
	}
}

// Picker 自定义负载均衡的地方
// 实现Picker google.golang.org/grpc/balancer 接口
type Picker struct {
	conns map[string][]*conn
	mutex sync.Mutex
	l     logger.GrpcLogger
}

func (p *Picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	nodes, lb, err := p.filteredNodes(info.Ctx)

	if err != nil {
		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
	}

	if len(nodes) == 0 {
		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
	} //没有节点
	var total int   //总权重
	var maxCC *conn //初始节点

	for _, cc := range nodes {
		total += cc.weight
		cc.currentWeight = cc.currentWeight + cc.weight
		if maxCC == nil || cc.currentWeight > maxCC.currentWeight {
			maxCC = cc
		}
	}
	maxCC.currentWeight = maxCC.currentWeight - total //更新权重
	p.l.Info("RPC调用", logger.Field{
		Key: "SWRR-节点选择",
		Val: fmt.Sprintf("%s该节点地址为[%s]", lb, maxCC.addr),
	})
	return balancer.PickResult{
		SubConn: maxCC.cc,
	}, nil
}

// filteredNodes 根据 ctx 获取的 label 来过滤对应节点 没有找到指定的节点给默认的(ordinary 普通节点)节点
func (p *Picker) filteredNodes(ctx context.Context) ([]*conn, string, error) {
	label := ctx.Value("label")
	lb, ok := label.(string)

	var builder strings.Builder
	builder.WriteString(fmt.Sprintf("当前用户为[%s]用户,", lb))

	if !ok {
		lb = defaultLabelType
	}
	notNode := true

	nodes, ok := p.conns[lb]
	if !ok || len(nodes) == 0 {
		notNode = false
	}
	if !notNode && lb != defaultLabelType {
		lb = defaultLabelType
		nodes, ok = p.conns[defaultLabelType]
	}
	if nodes == nil {
		return []*conn{}, "", balancer.ErrNoSubConnAvailable
	}
	builder.WriteString(fmt.Sprintf("选中[%s]节点,", lb))
	return nodes, builder.String(), nil

}

// 代表节点
type conn struct {
	weight        int    //真实权重
	currentWeight int    //当前权重
	addr          string //节点地址
	lb            string //节点类型
	cc            balancer.SubConn
}
